import cv2
import numpy as np
import torch


def preprocess_image(bgr_image: np.ndarray, size=(512, 1024)):
    h, w = bgr_image.shape[:2]
    rgb = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
    if size is not None:
        rgb = cv2.resize(rgb, (size[1], size[0]), interpolation=cv2.INTER_LINEAR)
    img = rgb.astype(np.float32) / 255.0
    img = (img - 0.5) / 0.5  # 简单归一化到 [-1,1]
    img = np.transpose(img, (2, 0, 1))  # CHW
    tensor = torch.from_numpy(img).unsqueeze(0).float()
    return tensor


def postprocess_output(logits: torch.Tensor, out_size):
    # logits: (1, C, H, W) 或 (C, H, W)
    if logits.dim() == 3:
        logits = logits.unsqueeze(0)
    pred = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.uint8)
    pred = cv2.resize(pred, (out_size[1], out_size[0]), interpolation=cv2.INTER_NEAREST)
    return pred


def apply_roi(mask: np.ndarray, top_ratio: float = 0.4):
    h, w = mask.shape[:2]
    roi_mask = np.zeros_like(mask)
    start_y = int(h * top_ratio)
    roi_mask[start_y:, :] = mask[start_y:, :]
    return roi_mask


def calculate_road_position(road_mask: np.ndarray):
    # 计算道路占比与相对中心偏移（-1~1）
    h, w = road_mask.shape[:2]
    ys, xs = np.where(road_mask > 0)
    total = len(xs)
    if total == 0:
        return 0.0, 0.0
    center_x = np.mean(xs)
    offset = (center_x - w / 2) / (w / 2)
    ratio = total / (h * w)
    return float(ratio), float(offset)